import torch
import models_xh as models
from data_512 import prepare_image_cv2
import cv2
import numpy as np


def cvt2vs(model, img):
    traced_script_module = torch.jit.trace(model, img)
    output = traced_script_module(img)
    traced_script_module.save(r"seg_model_jin.pt")
    result = output.squeeze().detach().cpu().numpy()
    result = (result * 255).astype(np.uint8)
    cv2.imshow("trace", cv2.resize(result, (600, 600)))


def pred(model, img):
    outs = model(img)
    result = outs.squeeze().detach().cpu().numpy()
    result = (result * 255).astype(np.uint8)
    cv2.imshow("pred",cv2.resize(result, (600, 600)))


if __name__ == '__main__':
    # device = "cpu"
    device = "cuda"
    model = models.resnet80().to(torch.device(device))
    model.load_state_dict(torch.load(r'./good/xhrs80_j512_184.pt', map_location=device))
    model.eval()
    print(model)
    img_u8 = cv2.imread("01_bingsi_4.jpg", 1)
    cv2.imshow("img_u8", cv2.resize(img_u8, (600, 600)))
    img_u8 = cv2.cvtColor(img_u8, cv2.COLOR_BGR2GRAY)
    equ = cv2.equalizeHist(img_u8)
    img_u8 = cv2.cvtColor(equ, cv2.COLOR_GRAY2BGR)
    original_img = np.array(img_u8, dtype=np.float32)
    h, w, _ = original_img.shape
    img = prepare_image_cv2(original_img)
    img = torch.from_numpy(img).unsqueeze(0).cuda()
    pred(model, img)
    cvt2vs(model, img)
    cv2.waitKey(0)
